# -*- coding: utf-8 -*-
"""
@Time ： 3/12/24 9:27 AM
@Auth ： woldier wong
@File ：det_utils.py
@IDE ：PyCharm
@DESCRIPTION：TODO
"""
import math
from collections import OrderedDict
from typing import Dict, List, Optional, Tuple

import torch
from torch import nn, Tensor
from torch.nn import functional as F
from torchvision.ops import complete_box_iou_loss, distance_box_iou_loss, FrozenBatchNorm2d, generalized_box_iou_loss


class BalancedPositiveNegativeSampler:
    """
    This class samples batches, ensuring that they contain a fixed proportion of positives
    本类采样batch, 确保负样本的个数是得到修正的
    """

    def __init__(self, batch_size_per_image: int, positive_fraction: float) -> None:
        """
        Args:
            batch_size_per_image (int): number of elements to be selected per image 一张图片中被选中的元素的个数
            positive_fraction (float): percentage of positive elements per batch  负样本所占总元素个数的比例
        """
        self.batch_size_per_image = batch_size_per_image
        self.positive_fraction = positive_fraction

    def __call__(self, matched_idxs: List[Tensor]) -> Tuple[List[Tensor], List[Tensor]]:
        """
        Args:
            matched idxs: list of tensors containing -1, 0 or positive values.
                Each tensor corresponds to a specific image.
                一个list 的tensor, 其中的值包括-1,0,或者正整数值, 每一个tensor 代表batch中的一个tensor
                -1 values are ignored, 0 are considered as negatives and > 0 as
                positives.
                -1 被忽略, 0被认为是负样本, >0 被认为是正样本
        Returns:
            pos_idx (list[tensor])  正样本的遮罩
            neg_idx (list[tensor])  负样本的遮罩

        Returns two lists of binary masks for each image.
        The first list contains the positive elements that were selected,
        and the second list the negative example.
        返回每张图片的binary mask(二进制遮罩).
        第一个list表示被选择的正样本遮罩
        第二个list表示被选择的负样本遮罩
        """
        pos_idx = []
        neg_idx = []
        for matched_idxs_per_image in matched_idxs:
            positive = torch.where(matched_idxs_per_image >= 1)[0]  # 正样本索引
            negative = torch.where(matched_idxs_per_image == 0)[0]  # 负样本索引

            num_pos = int(self.batch_size_per_image * self.positive_fraction)  # 计算正样本指标
            # protect against not enough positive examples
            num_pos = min(positive.numel(), num_pos)  # 与当前拥有的正样本指标 做比较, 取较小的一个
            num_neg = self.batch_size_per_image - num_pos  # 计算得到负样本数量
            # protect against not enough negative examples
            num_neg = min(negative.numel(), num_neg)

            # randomly select positive and negative examples
            perm1 = torch.randperm(positive.numel(), device=positive.device)[:num_pos]  # 打乱然后取出前num_pos个
            perm2 = torch.randperm(negative.numel(), device=negative.device)[:num_neg]# 打乱然后取出前num_neg个

            pos_idx_per_image = positive[perm1]  # 得到positive的索引
            neg_idx_per_image = negative[perm2]  # 得到neg的索引

            # create binary mask from indices
            pos_idx_per_image_mask = torch.zeros_like(matched_idxs_per_image, dtype=torch.uint8)  # 创建mask
            neg_idx_per_image_mask = torch.zeros_like(matched_idxs_per_image, dtype=torch.uint8)

            pos_idx_per_image_mask[pos_idx_per_image] = 1  # 将对应索引处赋值1
            neg_idx_per_image_mask[neg_idx_per_image] = 1

            pos_idx.append(pos_idx_per_image_mask)
            neg_idx.append(neg_idx_per_image_mask)

        return pos_idx, neg_idx


@torch.jit._script_if_tracing
def encode_boxes(reference_boxes: Tensor, proposals: Tensor, weights: Tensor) -> Tensor:
    """
    Encode a set of proposals with respect to some
    reference boxes

    Args:
        reference_boxes (Tensor): reference boxes
        proposals (Tensor): boxes to be encoded
        weights (Tensor[4]): the weights for ``(x, y, w, h)``
    """

    # perform some unpacking to make it JIT-fusion friendly
    wx = weights[0]
    wy = weights[1]
    ww = weights[2]
    wh = weights[3]

    proposals_x1 = proposals[:, 0].unsqueeze(1)
    proposals_y1 = proposals[:, 1].unsqueeze(1)
    proposals_x2 = proposals[:, 2].unsqueeze(1)
    proposals_y2 = proposals[:, 3].unsqueeze(1)

    reference_boxes_x1 = reference_boxes[:, 0].unsqueeze(1)
    reference_boxes_y1 = reference_boxes[:, 1].unsqueeze(1)
    reference_boxes_x2 = reference_boxes[:, 2].unsqueeze(1)
    reference_boxes_y2 = reference_boxes[:, 3].unsqueeze(1)

    # implementation starts here  计算 cx cy h w
    ex_widths = proposals_x2 - proposals_x1
    ex_heights = proposals_y2 - proposals_y1
    ex_ctr_x = proposals_x1 + 0.5 * ex_widths
    ex_ctr_y = proposals_y1 + 0.5 * ex_heights

    gt_widths = reference_boxes_x2 - reference_boxes_x1
    gt_heights = reference_boxes_y2 - reference_boxes_y1
    gt_ctr_x = reference_boxes_x1 + 0.5 * gt_widths
    gt_ctr_y = reference_boxes_y1 + 0.5 * gt_heights

    targets_dx = wx * (gt_ctr_x - ex_ctr_x) / ex_widths
    targets_dy = wy * (gt_ctr_y - ex_ctr_y) / ex_heights
    targets_dw = ww * torch.log(gt_widths / ex_widths)
    targets_dh = wh * torch.log(gt_heights / ex_heights)

    targets = torch.cat((targets_dx, targets_dy, targets_dw, targets_dh), dim=1)
    return targets


class BoxCoder:
    """
    This class encodes and decodes a set of bounding boxes into
    the representation used for training the regressors.

    本类编码和解码一批 bounding boxes 作为 用于训练回归预测的 表示
    """

    def __init__(
            self, weights: Tuple[float, float, float, float], bbox_xform_clip: float = math.log(1000.0 / 16)
    ) -> None:
        """
        Args:
            weights (4-element tuple)
            bbox_xform_clip (float)
        """
        self.weights = weights
        self.bbox_xform_clip = bbox_xform_clip

    def encode(self, reference_boxes: List[Tensor], proposals: List[Tensor]) -> List[Tensor]:
        """
        通过GT 和proposal 计算 回归参数
        :param reference_boxes:
        :param proposals:
        :return:
        """
        boxes_per_image = [len(b) for b in reference_boxes]  # 统计每张图片的anchor数目
        reference_boxes = torch.cat(reference_boxes, dim=0)  # 拼接
        proposals = torch.cat(proposals, dim=0)
        targets = self.encode_single(reference_boxes, proposals)
        return targets.split(boxes_per_image, 0)  # 拆分

    def encode_single(self, reference_boxes: Tensor, proposals: Tensor) -> Tensor:
        """
        Encode a set of proposals with respect to some
        reference boxes  计算回归参数

        Args:
            reference_boxes (Tensor): reference boxes
            proposals (Tensor): boxes to be encoded
        """
        dtype = reference_boxes.dtype
        device = reference_boxes.device
        weights = torch.as_tensor(self.weights, dtype=dtype, device=device)
        targets = encode_boxes(reference_boxes, proposals, weights)

        return targets

    def decode(self, rel_codes: Tensor, boxes: List[Tensor]) -> Tensor:
        """
        解码
        :param rel_codes: 回归预测参数 [B*A,4] B:batch_size ,A:一个图片的所有特征图中的anchors的合集
        :param boxes: 一批boxes(anchors) 一个list数组， list长度是batch_size, list中的一个tensor的shape[A, 4]
        :return:
        """
        torch._assert(
            isinstance(boxes, (list, tuple)),
            "This function expects boxes of type list or tuple.",
        )
        torch._assert(
            isinstance(rel_codes, torch.Tensor),
            "This function expects rel_codes of type torch.Tensor.",
        )
        # 得到每一张image 包含的anchor的数目
        boxes_per_image = [b.size(0) for b in boxes]
        # 将不同图片的anchor拼接在一起
        concat_boxes = torch.cat(boxes, dim=0)
        box_sum = 0  # 记录box 总数
        for val in boxes_per_image:
            box_sum += val
        if box_sum > 0:
            rel_codes = rel_codes.reshape(box_sum, -1)
        pred_boxes = self.decode_single(rel_codes, concat_boxes)
        if box_sum > 0:
            pred_boxes = pred_boxes.reshape(box_sum, -1, 4)
        return pred_boxes

    def decode_single(self, rel_codes: Tensor, boxes: Tensor) -> Tensor:
        """
        From a set of original boxes and encoded relative box offsets,
        get the decoded boxes.

        Args:
            rel_codes (Tensor): encoded boxes 回归预测参数,坐标偏移信息 [B*A,4] B:batch_size ,A:一个图片的所有特征图中的anchors的合集
            boxes (Tensor): reference boxes.  所有图片的anchors坐标信息 [B*A,4]
        """

        boxes = boxes.to(rel_codes.dtype)  # 将boxes的dtype 转换为与归回预测参数dtype相同

        widths = boxes[:, 2] - boxes[:, 0]  # 得到所有boxes的宽
        heights = boxes[:, 3] - boxes[:, 1]  # 得到所有boxes的高
        ctr_x = boxes[:, 0] + 0.5 * widths  # 得到中心x
        ctr_y = boxes[:, 1] + 0.5 * heights  # 得到中心y

        wx, wy, ww, wh = self.weights  # 得到权重
        dx = rel_codes[:, 0::4] / wx  # 这里rel_codes[:, 0::4] 才能保持原始的shape 如果用rel_codes[:, 0],最后一个维度便消失了
        dy = rel_codes[:, 1::4] / wy
        dw = rel_codes[:, 2::4] / ww
        dh = rel_codes[:, 3::4] / wh

        # Prevent sending too large values into torch.exp()
        dw = torch.clamp(dw, max=self.bbox_xform_clip)  # 限制范围上限  math.log(1000.0 / 16)
        dh = torch.clamp(dh, max=self.bbox_xform_clip)

        pred_ctr_x = dx * widths[:, None] + ctr_x[:, None]
        pred_ctr_y = dy * heights[:, None] + ctr_y[:, None]
        pred_w = torch.exp(dw) * widths[:, None]
        pred_h = torch.exp(dh) * heights[:, None]

        # Distance from center to box's corner.
        c_to_c_h = torch.tensor(0.5, dtype=pred_ctr_y.dtype, device=pred_h.device) * pred_h
        c_to_c_w = torch.tensor(0.5, dtype=pred_ctr_x.dtype, device=pred_w.device) * pred_w
        # 将 x y h w 转换为 x y x y
        pred_boxes1 = pred_ctr_x - c_to_c_w
        pred_boxes2 = pred_ctr_y - c_to_c_h
        pred_boxes3 = pred_ctr_x + c_to_c_w
        pred_boxes4 = pred_ctr_y + c_to_c_h
        # 拼接
        pred_boxes = torch.stack((pred_boxes1, pred_boxes2, pred_boxes3, pred_boxes4), dim=2).flatten(1)
        return pred_boxes


class BoxLinearCoder:
    """
    The linear box-to-box transform defined in FCOS. The transformation is parameterized
    by the distance from the center of (square) src box to 4 edges of the target box.
    """

    def __init__(self, normalize_by_size: bool = True) -> None:
        """
        Args:
            normalize_by_size (bool): normalize deltas by the size of src (anchor) boxes.
        """
        self.normalize_by_size = normalize_by_size

    def encode(self, reference_boxes: Tensor, proposals: Tensor) -> Tensor:
        """
        Encode a set of proposals with respect to some reference boxes

        Args:
            reference_boxes (Tensor): reference boxes
            proposals (Tensor): boxes to be encoded

        Returns:
            Tensor: the encoded relative box offsets that can be used to
            decode the boxes.

        """

        # get the center of reference_boxes
        reference_boxes_ctr_x = 0.5 * (reference_boxes[..., 0] + reference_boxes[..., 2])
        reference_boxes_ctr_y = 0.5 * (reference_boxes[..., 1] + reference_boxes[..., 3])

        # get box regression transformation deltas
        target_l = reference_boxes_ctr_x - proposals[..., 0]
        target_t = reference_boxes_ctr_y - proposals[..., 1]
        target_r = proposals[..., 2] - reference_boxes_ctr_x
        target_b = proposals[..., 3] - reference_boxes_ctr_y

        targets = torch.stack((target_l, target_t, target_r, target_b), dim=-1)

        if self.normalize_by_size:
            reference_boxes_w = reference_boxes[..., 2] - reference_boxes[..., 0]
            reference_boxes_h = reference_boxes[..., 3] - reference_boxes[..., 1]
            reference_boxes_size = torch.stack(
                (reference_boxes_w, reference_boxes_h, reference_boxes_w, reference_boxes_h), dim=-1
            )
            targets = targets / reference_boxes_size
        return targets

    def decode(self, rel_codes: Tensor, boxes: Tensor) -> Tensor:

        """
        From a set of original boxes and encoded relative box offsets,
        get the decoded boxes.

        Args:
            rel_codes (Tensor): encoded boxes
            boxes (Tensor): reference boxes.

        Returns:
            Tensor: the predicted boxes with the encoded relative box offsets.

        .. note::
            This method assumes that ``rel_codes`` and ``boxes`` have same size for 0th dimension. i.e. ``len(rel_codes) == len(boxes)``.

        """

        boxes = boxes.to(dtype=rel_codes.dtype)

        ctr_x = 0.5 * (boxes[..., 0] + boxes[..., 2])
        ctr_y = 0.5 * (boxes[..., 1] + boxes[..., 3])

        if self.normalize_by_size:
            boxes_w = boxes[..., 2] - boxes[..., 0]
            boxes_h = boxes[..., 3] - boxes[..., 1]

            list_box_size = torch.stack((boxes_w, boxes_h, boxes_w, boxes_h), dim=-1)
            rel_codes = rel_codes * list_box_size

        pred_boxes1 = ctr_x - rel_codes[..., 0]
        pred_boxes2 = ctr_y - rel_codes[..., 1]
        pred_boxes3 = ctr_x + rel_codes[..., 2]
        pred_boxes4 = ctr_y + rel_codes[..., 3]

        pred_boxes = torch.stack((pred_boxes1, pred_boxes2, pred_boxes3, pred_boxes4), dim=-1)
        return pred_boxes


class Matcher:
    """
    This class assigns to each predicted "element" (e.g., a box) a ground-truth
    element. Each predicted element will have exactly zero or one matches; each
    ground-truth element may be assigned to zero or more predicted elements.
    本class为每一个预测的element(例如box)分配一个GT element.
    每一个预测的elem将会拥有0个或者更多的GT-elem. 每个GT-elem 将会被分配0个或者更多的pre-elem

    Matching is based on the MxN match_quality_matrix, that characterizes how well
    each (ground-truth, predicted)-pair match. For example, if the elements are
    boxes, the matrix may contain box IoU overlap values.
    Matching 方法基于MxN match_quality_matrix(match 评估矩阵),
    其定义了两个match对(GT-elem, pre-elem)有多么的match.
    举个例子, 如果elem 是 boxes, 那么, 矩阵中保存的则是IoU.

    The matcher returns a tensor of size N containing the index of the ground-truth
    element m that matches to prediction n. If there is no match, a negative value
    is returned.
    matcher返回一个大小为N的tensor, 该tensor包含了GT-elem m 匹配的pre-elem n. 如果不存在匹配则会返回负数值.
    """

    BELOW_LOW_THRESHOLD = -1
    BETWEEN_THRESHOLDS = -2

    __annotations__ = {
        "BELOW_LOW_THRESHOLD": int,
        "BETWEEN_THRESHOLDS": int,
    }

    def __init__(self, high_threshold: float, low_threshold: float, allow_low_quality_matches: bool = False) -> None:
        """
        Args:
            high_threshold (float): quality values greater than or equal to
                this value are candidate matches.
                quality values大于等于本值的则是匹配候选elem
            low_threshold (float): a lower quality threshold used to stratify
                matches into three levels:
                 一个较低的阈值, 用于三个水平下满足匹配:
                1) matches >= high_threshold  满足条件时
                2) BETWEEN_THRESHOLDS matches in [low_threshold, high_threshold) H,L 之间的情况
                3) BELOW_LOW_THRESHOLD matches in [0, low_threshold) 低于L 的情况

            allow_low_quality_matches (bool): if True, produce additional matches
                for predictions that have only low-quality match candidates. See
                set_low_quality_matches_ for more details.
                如果给为True， 进行对pre（该pre只是低质量的候选elem）进行额外的match。
                通过 set_low_quality_matches_ function 查询更多的。
        """
        self.BELOW_LOW_THRESHOLD = -1
        self.BETWEEN_THRESHOLDS = -2
        torch._assert(low_threshold <= high_threshold, "low_threshold should be <= high_threshold")
        self.high_threshold = high_threshold
        self.low_threshold = low_threshold
        self.allow_low_quality_matches = allow_low_quality_matches

    def __call__(self, match_quality_matrix: Tensor) -> Tensor:
        """
        Args:
            match_quality_matrix (Tensor[float]): an MxN tensor, containing the
            pairwise quality between M ground-truth elements and N predicted elements.
            一个 MxN 张量，包含 M 个GT elem 和 N 个pre-elem之间的成对质量。
        Returns:
            matches (Tensor[int64]): an N tensor where N[i] is a matched gt in
            [0, M - 1] or a negative value indicating that prediction i could not
            be matched.
            返回一个长度为N的张量, N[i]的值是mtach的GT elem索引[0, M - 1],
            或者一个负值表明第i个预测, 没有匹配到GT elem.
        """
        if match_quality_matrix.numel() == 0:  # 判断GT 或者 pre 是否为空
            # empty targets or proposals not supported during training
            if match_quality_matrix.shape[0] == 0:
                raise ValueError("No ground-truth boxes available for one of the images during training")
            else:
                raise ValueError("No proposal boxes available for one of the images during training")

        # match_quality_matrix is M (gt) x N (predicted)
        # Max over gt elements (dim 0) to find best gt candidate for each prediction
        #    pre    0      1     2     3     4     5
        # GT \
        #  0       0.61  0.43   0.61  0.85  0.66  0.32
        #  1       0.81  0.57   0.61  0.45  0.18  0.66
        #  2       0.61  0.49   0.67  0.70  0.86  0.36
        # 当前pre和哪一个gt的匹配值最大， mactch 给出的是某个pre对应的最大gt的索引， matched_vals给出的是值。
        matched_vals, matches = match_quality_matrix.max(dim=0)  # 所有GT中匹配度最高的
        if self.allow_low_quality_matches:
            all_matches = matches.clone()  # 拷贝一份
        else:
            all_matches = None  # type: ignore[assignment]

        # Assign candidate matches with low quality to negative (unassigned) values
        # 对于每个elem匹配到的最大值， 判断是否是小于self.low_threshold，然后得到其mask
        below_low_threshold = matched_vals < self.low_threshold
        # 对与elem匹配的最大值，其值在low_threshold到high_threshold之间的mask
        between_thresholds = (matched_vals >= self.low_threshold) & (matched_vals < self.high_threshold)
        matches[below_low_threshold] = self.BELOW_LOW_THRESHOLD  # 设置匹配值在low_threshold之下的elem 的match 索引标记为-1
        matches[between_thresholds] = self.BETWEEN_THRESHOLDS # 设置匹配值在low_threshold和height_threshold的elem 的match 索引标记为-2

        if self.allow_low_quality_matches:
            if all_matches is None:
                torch._assert(False, "all_matches should not be None")
            else:
                # 这里matches 已经是经过阈值处理后的 有-1，-2索引，
                # 而all_match则没有经过阈值处理， 因此其索引值只能是大于等于0的
                self.set_low_quality_matches_(matches, all_matches, match_quality_matrix)

        return matches

    def set_low_quality_matches_(self, matches: Tensor, all_matches: Tensor, match_quality_matrix: Tensor) -> None:
        """
        Produce additional matches for predictions that have only low-quality matches.
        Specifically, for each ground-truth find the set of predictions that have
        maximum overlap with it (including ties); for each prediction in that set, if
        it is unmatched, then match it to the ground-truth with which it has the highest
        quality value.
        对于预测结果只有low-quality匹配的, 做额外的match.
        具体来说，对于每个GT ，找出与之有最大重合度（包括并列）的 pre;
        对于上述并列的pre, 如果没有匹配的, 将其与GT进行匹配.

        """
        # For each gt, find the prediction with which it has highest quality
        # 对于每一个gt 找到匹配值最大的pre-elem的值（highest_quality_foreach_gt），和索引（_）
        # 这里我们不需要pre-elem索引是因为有可能存在多个因此， 必须要根据每个gt的匹配最大值去找
        highest_quality_foreach_gt, _ = match_quality_matrix.max(dim=1)
        # Find highest quality match available, even if it is low, including ties
        # 原来的highest_quality_foreach_gt shape是[gt_num] 
        # 为了和matrix（shape[gt_num, pre_num]）做where查询，
        # 应该对highest_quality_foreach_gt做阔维 变成[gt_num,1], 这样就是到按照第0维找最大
        gt_pred_pairs_of_highest_quality = torch.where(match_quality_matrix == highest_quality_foreach_gt[:, None])
        # Example gt_pred_pairs_of_highest_quality:
        #   tensor([[    0, 39796],
        #           [    1, 32055],
        #           [    1, 32070],
        #           [    2, 39190],
        #           [    2, 40255],
        #           [    3, 40390],
        #           [    3, 41455],
        #           [    4, 45470],
        #           [    5, 45325],
        #           [    5, 46390]])
        # Each row is a (gt index, prediction index)
        # Note how gt items 1, 2, 3, and 5 each have two ties

        pred_inds_to_update = gt_pred_pairs_of_highest_quality[1]
        matches[pred_inds_to_update] = all_matches[pred_inds_to_update]


class SSDMatcher(Matcher):
    def __init__(self, threshold: float) -> None:
        super().__init__(threshold, threshold, allow_low_quality_matches=False)

    def __call__(self, match_quality_matrix: Tensor) -> Tensor:
        matches = super().__call__(match_quality_matrix)

        # For each gt, find the prediction with which it has the highest quality
        _, highest_quality_pred_foreach_gt = match_quality_matrix.max(dim=1)
        matches[highest_quality_pred_foreach_gt] = torch.arange(
            highest_quality_pred_foreach_gt.size(0), dtype=torch.int64, device=highest_quality_pred_foreach_gt.device
        )

        return matches


def overwrite_eps(model: nn.Module, eps: float) -> None:
    """
    This method overwrites the default eps values of all the
    FrozenBatchNorm2d layers of the model with the provided value.
    This is necessary to address the BC-breaking change introduced
    by the bug-fix at pytorch/vision#2933. The overwrite is applied
    only when the pretrained weights are loaded to maintain compatibility
    with previous versions.

    Args:
        model (nn.Module): The model on which we perform the overwrite.
        eps (float): The new value of eps.
    """
    for module in model.modules():
        if isinstance(module, FrozenBatchNorm2d):
            module.eps = eps


def retrieve_out_channels(model: nn.Module, size: Tuple[int, int]) -> List[int]:
    """
    This method retrieves the number of output channels of a specific model.

    Args:
        model (nn.Module): The model for which we estimate the out_channels.
            It should return a single Tensor or an OrderedDict[Tensor].
        size (Tuple[int, int]): The size (wxh) of the input.

    Returns:
        out_channels (List[int]): A list of the output channels of the model.
    """
    in_training = model.training
    model.eval()

    with torch.no_grad():
        # Use dummy data to retrieve the feature map sizes to avoid hard-coding their values
        device = next(model.parameters()).device
        tmp_img = torch.zeros((1, 3, size[1], size[0]), device=device)
        features = model(tmp_img)
        if isinstance(features, torch.Tensor):
            features = OrderedDict([("0", features)])
        out_channels = [x.size(1) for x in features.values()]

    if in_training:
        model.train()

    return out_channels


@torch.jit.unused
def _fake_cast_onnx(v: Tensor) -> int:
    return v  # type: ignore[return-value]


def _topk_min(input: Tensor, orig_kval: int, axis: int) -> int:
    """
    ONNX spec requires the k-value to be less than or equal to the number of inputs along
    provided dim. Certain models use the number of elements along a particular axis instead of K
    if K exceeds the number of elements along that axis. Previously, python's min() function was
    used to determine whether to use the provided k-value or the specified dim axis value.
    ONNX 规范要求 K 值小于或等于所提供维数的输入数。
    如果 K 超过了沿特定轴的元素数量，某些模型会使用该轴的元素数量而不是 K。
    以前，python 的 min() 函数用于确定使用提供的 K 值还是指定的 dim 轴值。

    However in cases where the model is being exported in tracing mode, python min() is
    static causing the model to be traced incorrectly and eventually fail at the topk node.
    In order to avoid this situation, in tracing mode, torch.min() is used instead.
    但是，在跟踪模式下导出模型时，python min() 是静态的，会导致模型被错误地跟踪，最终在 topk 节点失败。
    为了避免这种情况，在跟踪模式下，将使用 torch.min() 代替。

    Args:
        input (Tensor): The orignal input tensor. 最初的输入tensor
        orig_kval (int): The provided k-value.  k 值
        axis(int): Axis along which we retreive the input size.

    Returns:
        min_kval (int): Appropriately selected k-value.
    """
    if not torch.jit.is_tracing():  # 如果没有在tracing 直接调用py 自带的min 函数
        return min(orig_kval, input.size(axis))

    axis_dim_val = torch._shape_as_tensor(input)[axis].unsqueeze(0)
    min_kval = torch.min(torch.cat((torch.tensor([orig_kval], dtype=axis_dim_val.dtype), axis_dim_val), 0))
    return _fake_cast_onnx(min_kval)


def _box_loss(
        type: str,
        box_coder: BoxCoder,
        anchors_per_image: Tensor,
        matched_gt_boxes_per_image: Tensor,
        bbox_regression_per_image: Tensor,
        cnf: Optional[Dict[str, float]] = None,
) -> Tensor:
    torch._assert(type in ["l1", "smooth_l1", "ciou", "diou", "giou"], f"Unsupported loss: {type}")

    if type == "l1":
        target_regression = box_coder.encode_single(matched_gt_boxes_per_image, anchors_per_image)
        return F.l1_loss(bbox_regression_per_image, target_regression, reduction="sum")
    elif type == "smooth_l1":
        target_regression = box_coder.encode_single(matched_gt_boxes_per_image, anchors_per_image)
        beta = cnf["beta"] if cnf is not None and "beta" in cnf else 1.0
        return F.smooth_l1_loss(bbox_regression_per_image, target_regression, reduction="sum", beta=beta)
    else:
        bbox_per_image = box_coder.decode_single(bbox_regression_per_image, anchors_per_image)
        eps = cnf["eps"] if cnf is not None and "eps" in cnf else 1e-7
        if type == "ciou":
            return complete_box_iou_loss(bbox_per_image, matched_gt_boxes_per_image, reduction="sum", eps=eps)
        if type == "diou":
            return distance_box_iou_loss(bbox_per_image, matched_gt_boxes_per_image, reduction="sum", eps=eps)
        # otherwise giou
        return generalized_box_iou_loss(bbox_per_image, matched_gt_boxes_per_image, reduction="sum", eps=eps)


if __name__ == '__main__':
    pass
